require 'optparse'
require 'sqlite3'
require 'rubygems'
require 'active_record'
require 'mechanize'
require 'time'

$options = {}
parser = OptionParser.new("", 24) do |opts|
  opts.banner = "\nScraper 1.0\nAuthor: Louis (Skype: louisprm)\n\n"

  opts.on("-s", "--store PATH", "Path to output files") do |v|
    $options[:store] = v
  end

  opts.on("--syslog PATH", "Syslog file, default to /var/log/syslog") do |v|
    $options[:syslog] = v
  end

  opts.on("-i", "--interval SECOND", "Monitor interval in seconds, default to 60") do |v|
    $options[:interval] = v
  end

  opts.on_tail('-h', '--help', 'Displays this help') do
		puts opts
    exit
	end
end

parser.parse!

$options[:store] ||= File.dirname(__FILE__)

if !Dir.exists?($options[:store])
  raise "Directory #{$options[:store]} does not exists"
end

$options[:db] = File.join($options[:store], 'master.db')
$options[:log] = File.join($options[:store], 'httpguard.log')
$options[:last] = File.join($options[:store], '.last')

$options[:syslog] ||= '/var/log/syslog'
$options[:interval] ||= 60
$options[:interval] = $options[:interval].to_i

# Logger
$logger = Logger.new($options[:log])

ActiveRecord::Base.establish_connection(
  adapter: 'sqlite3',
  database: $options[:db],
  timeout: 15000
)

# Migration
class MySchema < ActiveRecord::Migration
  def change
    create_table :items do |t|
      t.string :ip
      t.integer :status
    end
  end
end

# initiate the database if not existed
MySchema.new.migrate(:change) unless File.exists?($options[:db])

class NSchedule
  def self.every(seconds)
    raise "No block provided" unless block_given?

    while true
      yield
      done_at = Time.now.to_i
      while Time.now.to_i - done_at < seconds.to_i
        sleep 1
      end
    end
  end  
end

class Item < ActiveRecord::Base
  NEW = 0
  BLOCKED = 1
  HTTP_PORTS = [80, 443]

  def self.last_checked
    Time.parse(IO.read($options[:last]))
  rescue Exception => ex
    nil
  end

  def self.update_last_checked(datetime)
    File.open($options[:last], 'w') {|f| f.write(datetime.to_s)}
    $logger.info "Last check time updated: #{datetime}"
  end

  def self.load
    last_check = self.last_checked
    datetime = nil

    $logger.info "Loading..."
    file = File.open($options[:syslog], 'r')
    ips = {}
    while !file.eof?
      line = file.readline
      datetime = Time.parse line[0..14]
      
      if last_check and last_check > datetime
        next
      end

      new_ip = line[/(?<=SRC=)[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}/] if line[/HTTP_FILTER[1-2]/]
      ips[new_ip] = true if new_ip
    end

    ips.keys.each do |new_ip|
      $logger.info "#{new_ip} detected"
      i = self.where(ip: new_ip).first_or_initialize
      i.status = NEW
      i.save
    end

    self.update_last_checked(datetime)
  end

  def self.load_unblock_and_block
    $logger.info "Starting..."
    self.load
     $logger.info "Checking..."
    self.all.each{ |ip|
      if ip.status == BLOCKED
        ip.unblock!
      else
        ip.block!
      end
    }
    $logger.info "Done!"
  end

  def block!
    HTTP_PORTS.each do |port|
      `sudo iptables -A INPUT -s #{self.ip} -p tcp --destination-port #{port} -j DROP`
    end
    
    self.status = BLOCKED
    self.save
    $logger.info "#{self.ip} blocked"
  end

  def unblock!
    #`sudo iptables -L --line-numbers | grep '#{self.ip}' | awk '{ print $1}' | xargs -I {} sudo iptables -D INPUT {}`
    ip = self.ip
    self.destroy
    $logger.info "#{ip} unblocked"
  end
end

# Run
NSchedule.every($options[:interval].to_i) do 
  Item.load_unblock_and_block
end